package net.majodali.util;

import java.util.Arrays;
import java.util.Iterator;
import java.util.Stack;

public class NestedIterator<T> implements Iterator<T> {
	Stack<Iterator<?>> iterators = new Stack<Iterator<?>>();
	T next;
	
	protected boolean isBranchNode(Object object) {
		return object != null && 
				(object instanceof Iterable<?> 
				|| object.getClass().isArray()); 
	}
	
	protected Iterator<?> getChildren(Object object) {
		if (object instanceof Iterable<?>)
			return ((Iterable) object).iterator();
		else if (object instanceof Object[])
			return Arrays.asList((Object[])object).iterator();
		else
			throw new IllegalArgumentException(String.valueOf(object));
	}
	
	private T getNext() {
		while (true) {
			if (iterators.isEmpty())
				return null;
			if (!iterators.peek().hasNext()) 
				iterators.pop();
			else {
				Object candidate = iterators.peek().next();
				if (isBranchNode(candidate))
					iterators.push(getChildren(candidate));
				else
					return (T) candidate;
			}
		}
	}
	
	public NestedIterator(Iterator<?> topLevelIterator) {
		iterators.push(topLevelIterator);
		next = getNext();
	}
	
	public NestedIterator(Iterable<?> source) {
		this(source.iterator());
	}
	
	@Override
	public boolean hasNext() {
		return !iterators.isEmpty();
	}

	@Override
	public T next() {
		return next;
	}

	@Override
	public void remove() {
		iterators.peek().remove();
	}
}
